package com.gitee.cirnochat.chatai.client;

import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;

import java.util.List;
import java.util.concurrent.Semaphore;

@RequiredArgsConstructor
public class SparkChatModel implements ChatModel {

    private final String baseUrl;

    private final String apiKey;

    private final String model;

    private final static int MAX_QPS = 2; // 每秒请求数

    private final static Semaphore clients = new Semaphore(MAX_QPS);

    @Override
    public ChatResponse call(Prompt prompt) {
        // 0. 获取信号量
        try {
            clients.acquire();
            // 1. 获取用户消息
            List<Message> messages = prompt.getInstructions();
            JSONArray messageArray = new JSONArray();

            for (Message msg : messages) {
                JSONObject msgJson = new JSONObject();
                msgJson.set("role", msg.getMessageType().getValue());
                msgJson.set("content", msg.getText());
                messageArray.add(msgJson);
            }

            // 2. 构建请求体
            JSONObject requestBody = new JSONObject();
            requestBody.set("model", model);
            requestBody.set("stream", false); // 当前实现非流式
            requestBody.set("messages", messageArray);

            // 3. 构建请求
            HttpRequest request = HttpUtil.createPost(baseUrl)
                    .header("Authorization", "Bearer " + apiKey)
                    .header("Content-Type", "application/json")
                    .body(requestBody.toString());

            try (HttpResponse response = request.execute()) {
                if (response.getStatus() == 200) {
                    JSONObject responseJson = new JSONObject(response.body());
                    JSONArray choices = responseJson.getJSONArray("choices");
                    if (!choices.isEmpty()) {
                        JSONObject firstChoice = choices.getJSONObject(0);
                        String content = firstChoice.getJSONObject("message").getStr("content");
                        return new ChatResponse(List.of(new Generation(new AssistantMessage(content))));
                    }
                }
                throw new RuntimeException("Spark API 请求失败：" + response.getStatus() + " - " + response.body());
            }
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } finally {
            clients.release();
        }
    }

    // 其他方法可以返回空或抛出异常视需求而定
    @Override
    public String call(String message) {
        return call(new Prompt(List.of(new UserMessage(message)))).getResult().getOutput().getText();
    }


}
